import os
import time
import re
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from utils.title_level1 import title_change
from utils.table_merge import table_change

start_time = time.time()
# 参数设置:文件路径、文件名
pdf_file_name = "pdf_files/03苏州国际发展集团有限公司2021度第一期中期票据募集说明书.pdf"  # 替换真实的pdf路径
# pdf_file_name = "/home/fengjie/MinerU/demo/pdf_files/03苏州国际发展集团有限公司2021度第一期中期票据募集说明书.pdf"  # 替换真实的pdf路径

name_without_suff = pdf_file_name.split(".")[0]

# 准备环境
local_image_dir, local_md_dir = "output/images", "output"
image_dir = str(os.path.basename(local_image_dir)) # image_dir = images
## 创建存储文件夹
os.makedirs(local_image_dir, exist_ok=True)
## 实例化writer对象：用于把数据写入文件中
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(
    local_md_dir
)

# read bytes
reader1 = FileBasedDataReader("")
pdf_bytes = reader1.read(pdf_file_name)  # 准备pdf目录

# 处理过程
## 创造数据集实例
ds = PymuDocDataset(pdf_bytes)

# 管道处理
# pipe_result = ds.apply(doc_analyze, ocr=False).pipe_txt_mode(image_writer)
pipe_result = ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer)

# 转换为markdown语法
md_content = pipe_result.get_markdown(image_dir)

# 标题再处理
md_content = title_change(md_content)
# 表格再处理
# md_content = table_change(md_content)

# 保存到markdown文件
md_writer.write_string(f"{name_without_suff}.md",md_content)

end_time = time.time()

print('总时间：',end_time-start_time)

